import { AgentMakerEnum } from "../../../public/agentBackend";
import { EmbeddingBackend, EmbeddingConfig } from "../embedding";
import { RETRIEVAL } from "./config";
import { MemoryStore, JsonMemoryStore } from "./memoryStore";

export interface AllMemoryConfig {
    /** 模型制造商  */
    modelMaker?: AgentMakerEnum;
    /** Embedding模型类型 */
    embeddingModelType?: string;
    /** 代码检索的topK */
    topKCode?: number;
    /** 文本检索的topK */
    topKText?: number;
    /** 代码检索的阈值 */
    codeThreshold?: number;
    /** 文本检索的阈值 */
    textThreshold?: number;
    /** embedding维度 */
    embeddingDimension?: number;
    /** 记忆存储对象 */
    memoryStore?: MemoryStore<string>;
    /** 文本embedding方法 */
    embeddingConfig?: EmbeddingConfig;
}

export class AllMemory {
    public modelMaker: AgentMakerEnum;
    public embeddingModelType: string;
    public topKCode: number;
    public topKText: number;
    public codeThreshold: number;
    public textThreshold: number;
    public embeddingDimension: number;
    public memoryStore: MemoryStore<string>;
    public embeddingConfig: EmbeddingConfig;

    /** embedding后端 */
    public embeddingBackend: EmbeddingBackend;

    constructor(config?: AllMemoryConfig) {
        this.modelMaker = config?.modelMaker || AgentMakerEnum.OpenAi;
        this.embeddingModelType = config?.embeddingModelType || "text-embedding-ada-002";
        this.topKCode = config?.topKCode || RETRIEVAL.TOP_K_CODE;
        this.topKText = config?.topKText || RETRIEVAL.TOP_K_TEXT;
        this.codeThreshold = config?.codeThreshold || RETRIEVAL.SEARCH_CODE_THRESHOLD;
        this.textThreshold = config?.textThreshold || RETRIEVAL.SEARCH_TEXT_THRESHOLD;
        this.embeddingDimension = config?.embeddingDimension || 768;
        this.memoryStore = config?.memoryStore || new JsonMemoryStore(this.embeddingDimension);
        this.embeddingConfig = config?.embeddingConfig || {
            modelMaker: this.modelMaker,
            embeddingModelType: this.embeddingModelType
        };
        this.embeddingBackend = new EmbeddingBackend(this.embeddingConfig);
    }

    loadMemory() {}
    async searchText(query: string, theK?: number, theThreshold?: number): Promise<string[]> {
        const k = theK ?? this.topKText;
        const threshold = theThreshold ?? this.textThreshold;
        if ((await this.getMemoryCount()) === 0 || !query || k <= 0) {
            return [];
        }
        const embeddingQuery = await this.embeddingBackend.getEmbedding(query);
        if (embeddingQuery.length < 0) {
            return [];
        }
        const results = await this.memoryStore.query(embeddingQuery, threshold, k);
        return results.map((item) => item.content);
    }
    async searcCode(query: string, theK?: number, theThreshold?: number): Promise<string[]> {
        const k = theK ?? this.topKText;
        const threshold = theThreshold ?? this.textThreshold;
        if ((await this.getMemoryCount()) === 0 || !query || k <= 0) {
            return [];
        }
        const embeddingQuery = await this.embeddingBackend.getEmbedding(query);
        if (embeddingQuery.length < 0) {
            return [];
        }
        const results = await this.memoryStore.query(embeddingQuery, threshold, k);
        return results.map((item) => item.content);
    }
    async retrieve(query: string, type: string, k?: number): Promise<string[]> {
        if (type === "text") {
            return await this.searchText(query, k ?? this.topKText);
        }
        if (type === "code") {
            return await this.searcCode(query, k ?? this.topKCode);
        }
        return [];
    }
    async memorize(content: string, type: string): Promise<boolean> {
        const embeddingContent = await this.embeddingBackend.getEmbedding(content);
        if (embeddingContent.length < 0) {
            return false;
        }
        return await this.memoryStore.save({ content, embedding: embeddingContent });
    }

    private async getMemoryCount(): Promise<number> {
        const stats = await this.memoryStore.getStatistics();
        return stats.size;
    }
}
